import networkx as nx
import random
import matplotlib.pyplot as plt
import community
import numpy as np
import pandas as pd
import seaborn as sns
from scipy import integrate
import scipy.linalg as sl


# 读取csv文件中的邻接矩阵，并根据邻接矩阵生成Graph
def get_csv_matrix(filename):
    df = pd.read_csv(filename, header=None, index_col=None)
    DiG = nx.from_pandas_adjacency(df, create_using=nx.DiGraph())
    return DiG
    pass

#   广义特征值与广义特征向量
def get_generalized_eigenvalue(origin_G, mst_G):
    # 计算拉普拉斯矩阵并转换为密集矩阵
    laplac_matrix_origin = nx.directed_laplacian_matrix(origin_G)
    laplac_matrix_mst = nx.directed_laplacian_matrix(mst_G)
    if mst_G.number_of_edges() == 1554:
        np.savetxt("origin_mat.csv", laplac_matrix_origin, delimiter=",")
        np.savetxt("mst_mat.csv", laplac_matrix_mst, delimiter=",")
    [e, v] = sl.eig(laplac_matrix_origin, laplac_matrix_mst)
    # e包含特征值，v包含对应的特征向量
    e.sort()
    # print("特征值:", e)
    return e
    # print("特征向量:\n", v)
    pass

def disparity_filter(G, weight='weight'):
    ''' Compute significance scores (alpha) for weighted edges in G as defined in Serrano et al. 2009
        Args
            G: Weighted NetworkX graph
        Returns
            Weighted graph with a significance score (alpha) assigned to each edge
        References
            M. A. Serrano et al. (2009) Extracting the Multiscale backbone of complex weighted networks. PNAS, 106:16, pp. 6483-6488.
    '''

    if nx.is_directed(G):  # directed case
        N = nx.DiGraph()
        for u in G:

            k_out = G.out_degree(u)
            k_in = G.in_degree(u)

            if k_out > 1:
                sum_w_out = sum(np.absolute(G[u][v][weight]) for v in G.successors(u))
                for v in G.successors(u):
                    w = G[u][v][weight]
                    p_ij_out = float(np.absolute(w)) / sum_w_out
                    alpha_ij_out = 1 - (k_out - 1) * integrate.quad(lambda x: (1 - x) ** (k_out - 2), 0, p_ij_out)[0]
                    N.add_edge(u, v, weight=w, alpha_out=float('%.4f' % alpha_ij_out))

            elif k_out == 1 and G.in_degree(G.successors(u)[0]) == 1:
                # we need to keep the connection as it is the only way to maintain the connectivity of the network
                v = G.successors(u)[0]
                w = G[u][v][weight]
                N.add_edge(u, v, weight=w, alpha_out=0., alpha_in=0.)
                # there is no need to do the same for the k_in, since the link is built already from the tail

            if k_in > 1:
                sum_w_in = sum(np.absolute(G[v][u][weight]) for v in G.predecessors(u))
                for v in G.predecessors(u):
                    w = G[v][u][weight]
                    p_ij_in = float(np.absolute(w)) / sum_w_in
                    alpha_ij_in = 1 - (k_in - 1) * integrate.quad(lambda x: (1 - x) ** (k_in - 2), 0, p_ij_in)[0]
                    N.add_edge(v, u, weight=w, alpha_in=float('%.4f' % alpha_ij_in))
        return N

    else:  # undirected case
        B = nx.Graph()
        for u in G:
            k = len(G[u])
            if k > 1:
                sum_w = sum(np.absolute(G[u][v][weight]) for v in G[u])
                for v in G[u]:
                    w = G[u][v][weight]
                    p_ij = float(np.absolute(w)) / sum_w
                    alpha_ij = 1 - (k - 1) * integrate.quad(lambda x: (1 - x) ** (k - 2), 0, p_ij)[0]
                    B.add_edge(u, v, weight=w, alpha=float('%.4f' % alpha_ij))
        return B


def disparity_filter_alpha_cut(G, weight='weight', alpha_t=0.4, cut_mode='or'):
    ''' Performs a cut of the graph previously filtered through the disparity_filter function.

        Args
        ----
        G: Weighted NetworkX graph

        weight: string (default='weight')
            Key for edge data used as the edge weight w_ij.

        alpha_t: double (default='0.4')
            The threshold for the alpha parameter that is used to select the surviving edges.
            It has to be a number between 0 and 1.

        cut_mode: string (default='or')
            Possible strings: 'or', 'and'.
            It works only for directed graphs. It represents the logic operation to filter out edges
            that do not pass the threshold value, combining the alpha_in and alpha_out attributes
            resulting from the disparity_filter function.


        Returns
        -------
        B: Weighted NetworkX graph
            The resulting graph contains only edges that survived from the filtering with the alpha_t threshold

        References
        ---------
        .. M. A. Serrano et al. (2009) Extracting the Multiscale backbone of complex weighted networks. PNAS, 106:16, pp. 6483-6488.
    '''

    if nx.is_directed(G):  # Directed case:
        B = nx.DiGraph()
        for u, v, w in G.edges(data=True):
            try:
                alpha_in = w['alpha_in']
            except KeyError:  # there is no alpha_in, so we assign 1. It will never pass the cut
                alpha_in = 1
            try:
                alpha_out = w['alpha_out']
            except KeyError:  # there is no alpha_out, so we assign 1. It will never pass the cut
                alpha_out = 1

            if cut_mode == 'or':
                if alpha_in < alpha_t or alpha_out < alpha_t:
                    B.add_edge(u, v, weight=w[weight])
            elif cut_mode == 'and':
                if alpha_in < alpha_t and alpha_out < alpha_t:
                    B.add_edge(u, v, weight=w[weight])
        return B

    else:
        B = nx.Graph()  # Undirected case:
        for u, v, w in G.edges(data=True):

            try:
                alpha = w['alpha']
            except KeyError:  # there is no alpha, so we assign 1. It will never pass the cut
                alpha = 1

            if alpha < alpha_t:
                B.add_edge(u, v, weight=w[weight])
        return B


def get_iqr_data(datas):
    q1 = np.quantile(datas, 0.25)
    q2 = np.median(datas)
    q3 = np.quantile(datas, 0.75)
    iqr = q3 - q1
    down = q1 - 1.5 * iqr
    up = q3 + 1.5 * iqr
    return [q1, q2, q3, iqr, down, up]

original_DiG = get_csv_matrix('wajueji.csv')
G = disparity_filter(G=original_DiG)
print('original: nodes = %s, edges = %s' % (G.number_of_nodes(), G.number_of_edges()))  # 打印原始图数据

alpha_in_Q1 = []
alpha_out_Q1 = []
for u, v, d in G.edges(data=True):
    alpha_in = d['alpha_in']  # 从边的属性中获取权重
    alpha_out = d['alpha_out']  # 从边的属性中获取权重
    alpha_in_Q1.append(alpha_in)
    alpha_out_Q1.append(alpha_out)
    pass
# 打印四分数前25%
print('weighted in-degree: Q1_in = %s, weighted out-degree = %s' % (get_iqr_data(alpha_in_Q1), get_iqr_data(alpha_out_Q1)))

DiMST = nx.tree.minimum_spanning_arborescence(G, "alpha_in")
print('Before: Edges of DiMST:  %s' % DiMST.number_of_edges())



x = []
y = []
old_num_edges = 0
# 向mst中添加边
for s in range(0, 100):
    for u, v, d in G.edges(data=True):
        if d['alpha_in'] < float(s)/100:
            DiMST.add_edge(u, v, **d)
        pass
    if DiMST.number_of_edges() == old_num_edges:
        continue
    # nx.write_edgelist(DiMST, "test-edgelist.csv", data=True)    # 生成edgelist
    print(' After: Edges of DiMST:  %s' % DiMST.number_of_edges())
    eigenval = get_generalized_eigenvalue(G, DiMST)
    if eigenval[len(eigenval) - 1].real > 100000:
        print('相对条件数: k = %s' % (eigenval[len(eigenval) - 2] / eigenval[1]))
        y.append(eigenval[len(eigenval) - 2] / eigenval[1])
    else:
        print('相对条件数: k = %s' % (eigenval[len(eigenval) - 1] / eigenval[0]))
        y.append(eigenval[len(eigenval) - 1] / eigenval[0])
    x.append(DiMST.number_of_edges())
    # print('相对条件数: k = %s' % (eigenval[len(eigenval) - 2]/eigenval[1]))

plt.plot(x, y)
plt.show()
plt.savefig("a.png")
# 创建一个DataFrame，索引和列都是节点列表
# 初始值都设为0，表示没有边
# nodes = list(DiMST.nodes())
# adj_matrix = pd.DataFrame(0, index=nodes, columns=nodes)
#
# # 遍历图的边
# for u, v, w in DiMST.edges(data=True):
#     adj_matrix.at[u, v] = w['weight']
#     if not DiMST.is_directed():  # 如果图是无向的，还需要设置反向的边
#         adj_matrix.at[v, u] = w['weight']
#
#     # 将邻接矩阵保存到CSV文件
# adj_matrix.to_csv('adjmatrix.csv')\
# 获取所有节点并按标签排序
# sorted_nodes = sorted(DiMST.nodes())

# 创建一个映射，将排序后的节点标签映射到索引上
# node_index_map = {node: idx for idx, node in enumerate(sorted_nodes)}

# 使用numpy的fancy indexing来重新排序邻接矩阵的行和列
# adj_matrix = nx.adjacency_matrix(DiMST, nodelist=sorted_nodes).todense()

# 使用pandas创建DataFrame，并指定索引和列
# df = pd.DataFrame(adj_matrix, index=sorted_nodes, columns=sorted_nodes)

# 将DataFrame导出为CSV文件
# df.to_csv('adjacency_matrix_sorted.csv')
# Press the green button in the gutter to run the script.
# if __name__ == '__main__':
#     print_hi('PyCharm')

# See PyCharm help at https://www.jetbrains.com/help/pycharm/
